import cv2
import argparse
import  os
import time
from sift import getkp,feature_inbox
# from light import getkp,feature_inbox
from MDFlow_main.demo_md import flowtrack,line_track
# from RAFT.demo import flowtrack,line_track
import  key_frame
import patch_track
import numpy as np

def parse_args():
    parser = argparse.ArgumentParser(description='track Demo')
    #huaxi
    # no_instru
    parser.add_argument('--images1', default='./rectified_left/', help='video1 path')
    parser.add_argument('--linepoint', default="/home/ltt/point_track_final/1_22.txt", help='linepoint path')
    parser.add_argument('--kf_num', default=3, help='kf_num')
    parser.add_argument('--box_side', default=0.35, help='box_side')
    parser.add_argument('--center', default=10, help='center_region')
    parser.add_argument('--time1', default=10, help='center_region')#相似关键帧>=2,时间间隔
    parser.add_argument('--time2', default=4, help='center_region')#相似关键帧=1,时间间隔
    args = parser.parse_args()
    return args

def line_trackpoint(data,img):
    for j in range(len(data)):
        cv2.circle(img, (int(data[j][0]), int(data[j][1])), 1, (0,  0,255), -1)
    return img

def get_kfram(args):
    k=args.kf_num#相似帧数
    dirs_init1 = os.listdir(args.images1)
    dirs_init1.sort(key=lambda x: int(x[:-4]))
    t_data = np.genfromtxt(args.linepoint, dtype=['float64', 'float64'], delimiter='')
    data = t_data.view(np.float64).reshape(-1, 2)
    data=np.unique(data,axis = 0)#剔除重复值
    data=key_frame.b_line(data)
    point_num = len(data)
    scale_threshold=[12,-12]
    kframe=[]#存关键帧

    img0= cv2.imread(args.images1 + '/{}'.format(dirs_init1[0]), flags=1)
    w = img0.shape[1]
    h=img0.shape[0]
    l =int(min(w,h)*args.box_side)

    w_min=args.center#判断路径中心是否在一定区域内，判断关键帧的选取
    w_max=w-w_min
    h_min=w_min
    h_max=h-h_min

    keypool_update = key_frame.keypool_update_factory()# 使用工厂函数创建 keypool_update 函数实例

    for i in range (len(dirs_init1)):
        print("计算到第{}帧".format(i))
        k_frame=key_frame.keyframe()

        img1_rgb= cv2.imread(args.images1 + '/{}'.format(dirs_init1[i]), flags=1)
        img1_rgb_copy=img1_rgb.copy()
        kp1,des1=getkp(img1_rgb)
        if i==0:#第0帧作为关键帧
            box_center_x,box_center_y= patch_track.box_center(data)#标记点的中心点
            box1=patch_track.boxin(box_center_x,box_center_y,l,img1_rgb)#矩形框
            kp1_in ,des1_in=feature_inbox(box1,kp1,des1)

            k_frame.index=i
            k_frame.imgrgb=img1_rgb_copy
            k_frame.box=box1
            k_frame.kp=kp1_in#只存在框内的点
            k_frame.des=des1_in
            k_frame.linepoint=data
            k_frame.pointcenter = np.array([box_center_x,box_center_y])
            k_frame.kpall=kp1
            k_frame.desall=des1

            kframe.append(k_frame)
            p=key_frame.b_line(k_frame.linepoint)
            img_box = key_frame.line_trackpoint(img1_rgb, p)


        elif i in range(1,k) :#后k-1帧直接使用光流，构成初始关键帧池
            u_f, v_f,_,_ = flowtrack(kframe[0].imgrgb, img1_rgb)
            data_new, img_draw, data_newin = line_track(u_f, v_f, kframe[0].linepoint, img1_rgb)#光流处
            box_center_x, box_center_y =patch_track.box_center(data_new)
            box1=patch_track.boxin(box_center_x, box_center_y,l,img1_rgb)

            kp1_in, des1_in = feature_inbox(box1, kp1, des1)
            k_frame.box=box1
            k_frame.linepoint=data_new
            k_frame.index = i
            k_frame.imgrgb = img1_rgb_copy
            k_frame.kp = kp1_in
            k_frame.des = des1_in
            k_frame.pointcenter = np.array([box_center_x, box_center_y])
            k_frame.kpall = kp1
            k_frame.desall = des1

            kframe.append(k_frame)
            img_line=key_frame.line_trackpoint(img1_rgb,k_frame.linepoint)
            img_box=img_line

        else:
            num_past=kframe[len(kframe)-1].index#最近一帧的索引
            simi_frame,simi_frame_better,simi_frame_nobetter=patch_track.traversal_kframe(kframe, kp1, des1, k)#找相似关键帧(kframe[m], len(goodin), [i[0] for i in goodin], [i[1] for i in goodin],good1,good2)
            if len(simi_frame_better)>=2:
                # print("相似关键帧为",simi_frame[0][0].index)
                track_final, data_scale,img_mask,rot= patch_track.keyframe_affine(img1_rgb, simi_frame_better, i,scale_threshold)
                if 2 < data_scale[0] < data_scale[1] < w - 2 and 2 < data_scale[2] < data_scale[3] < h - 2 :#关键帧中的线距离图片边界超过两个像素
                    y=1
                else:
                    y=0
                if len(track_final) == point_num:#追踪的点都在图片范围内
                    if i-num_past>=args.time1 and y==1:#间隔10帧存为关键帧
                        box_center_x, box_center_y = patch_track.box_center(track_final)
                        box1 = patch_track.boxin(box_center_x, box_center_y, l, img1_rgb)

                        kp1_in, des1_in = feature_inbox(box1, kp1, des1)
                        k_frame.index = i
                        k_frame.imgrgb = img1_rgb_copy
                        k_frame.box = box1
                        k_frame.kp = kp1_in
                        k_frame.des = des1_in
                        k_frame.linepoint = track_final
                        k_frame.pointcenter = np.array([box_center_x, box_center_y])
                        k_frame.kpall = kp1
                        k_frame.desall = des1

                        kframe =keypool_update(kframe, k_frame, simi_frame_better[0],rot)
                        track_final_ = key_frame.b_line(track_final)#插值
                        img_line = key_frame.line_trackpoint(img1_rgb, track_final_)
                        # img_box = key_frame.keyframe.line_box(k_frame, img_line)
                        # print("第{}帧为间隔的关键帧".format(i))
                        img_box=img_line
                    else:
                        track_final_ = key_frame.b_line(track_final)
                        img_line = key_frame.line_trackpoint(img1_rgb, track_final_)
                        img_box=img_line
                        # print("第{}_帧为普通帧".format(i))
                elif track_final.shape[0] == 0:#追踪的点全都不在图片上
                    img_box =img1_rgb_copy
                    # print("第{}帧为彻底的无效帧1".format(i))
                else:#追踪的点不全
                    track_final_ = key_frame.b_line(track_final)
                    img_line = key_frame.line_trackpoint(img1_rgb, track_final_)
                    img_box =img_line
                    # print("第{}_帧为追踪不全的帧".format(i))

            elif len(simi_frame_better)==0:
                img_box = img1_rgb_copy

            else :#len(simi_frame_better)==1
                track_final,data_scale,img_mask,rot= patch_track.keyframe_affine(img1_rgb, simi_frame_better,i,scale_threshold)
                track_final_ = key_frame.b_line(track_final)
                if 2 < data_scale[0] < data_scale[1] < w - 2 and 2 < data_scale[2] < data_scale[3] < h - 2 :
                    y=1
                else:
                    y=0
                if len(track_final) == point_num:
                    img_line = key_frame.line_trackpoint(img1_rgb, track_final_)  # 画线
                    if  i - num_past > args.time2 and y==1:
                        box_center_x, box_center_y = patch_track.box_center(track_final)
                        if w_min < box_center_x < w_max and h_min < box_center_y < h_max:
                            box1 = patch_track.boxin(box_center_x, box_center_y, l, img1_rgb)
                            kp1_in, des1_in = feature_inbox(box1, kp1, des1)
                            k_frame.index = i
                            k_frame.imgrgb = img1_rgb_copy
                            k_frame.box = box1
                            k_frame.kp = kp1_in
                            k_frame.des = des1_in
                            k_frame.linepoint = track_final
                            k_frame.pointcenter = np.array([box_center_x, box_center_y])
                            k_frame.kpall = kp1
                            k_frame.desall = des1

                            kframe = keypool_update(kframe, k_frame, simi_frame_better[0],rot)
                            img_box=img_line
                            # print("第{}帧为间隔的关键帧".format(i))
                        else:
                            img_box = img_line
                            # print("第{}帧为__普通帧".format(i))
                    else:
                        img_box = img_line
                        # print("第{}帧为普通帧".format(i))
                elif len(track_final)==0:
                    img_box=img1_rgb_copy
                    # print("第{}帧为彻底无效帧5".format(i))
                else:
                    img_box= key_frame.line_trackpoint(img1_rgb, track_final_)#画线_green
                    # print("第{}帧为追踪不全帧".format(i))
        # cv2.imwrite(os.path.join( track_img, f'{i:05}.jpg'), img_box)
        cv2.namedWindow("0", cv2.WINDOW_NORMAL)
        cv2.imshow("0", img_box)
        cv2.waitKey(1)
    return kframe


if __name__ == '__main__':
    args = parse_args()
    get_kfram(args)


